from docx import Document
from docx.oxml.ns import qn
import os
from pathlib import Path


def load_employee_file(file_path, image_save_dir="./employee_manual_images"):
    """
    仅针对 employee_manual.docx 读取：修复XPath表达式错误，支持段落、表格、内嵌图片（如组织架构图）
    :param file_path: employee_manual.docx 文件路径
    :param image_save_dir: 图片保存目录（默认：./employee_manual_images）
    :return: 整合所有内容的字符串
    """
    # 1. 初始化文档与结果存储（按原文档顺序保存，适配 employee_manual 的章节结构）
    doc = Document(file_path)
    all_content = []
    image_count = 0  # 图片序号（仅用于文档中唯一的组织架构图）
    file_name = Path(file_path).stem  # 提取文件名：employee_manual

    # 2. 创建图片保存目录（确保能存储组织架构图）
    os.makedirs(image_save_dir, exist_ok=True)

    # 3. 遍历文档 body 元素（按顺序处理段落、表格，匹配文档结构）
    for child in doc.element.body:
        # ------------------- 处理段落（含文本+内嵌图片，重点适配组织架构图段落）-------------------
        if child.tag == qn("w:p"):
            # 匹配当前段落对象（确保对应文档中的真实段落，如“2.1 公司组织架构图”段落）
            para_idx = len(doc.paragraphs) - len(doc.element.body.findall(qn("w:p"))) + doc.element.body.findall(
                qn("w:p")).index(child)
            para = doc.paragraphs[para_idx]
            para_text = []

            # 遍历段落中的文本块（run），提取文本和图片
            for run in para.runs:
                # 1）提取段落文本（如章节标题、说明文字）
                if run.text.strip():
                    para_text.append(run.text.strip())

                # 2）提取内嵌图片（关键修复：用正确语法定位 <a:blip>，适配组织架构图）
                # 方案：先获取 run 元素的所有子元素，筛选出 <a:blip> 标签
                for elem in run.element.iterchildren():
                    # 递归检查当前元素及子元素是否包含 <a:blip>
                    def find_blip(element):
                        if element.tag == qn("a:blip"):
                            return element
                        for child_elem in element.iterchildren():
                            result = find_blip(child_elem)
                            if result:
                                return result
                        return None

                    blip = find_blip(elem)
                    if blip:
                        rId = blip.get(qn("a:embed"))  # 获取图片资源ID
                        if rId:
                            # 从文档中获取图片资源（适配文档的图片存储格式）
                            image_part = doc.part.related_parts[rId]
                            # 识别图片格式（如 png/jpg，适配组织架构图）
                            content_type = image_part.content_type
                            img_ext = content_type.split("/")[-1] if "/" in content_type else "png"
                            # 图片保存路径（避免重名，标注文档来源）
                            img_save_path = os.path.join(image_save_dir, f"img_{image_count}_{file_name}.{img_ext}")

                            # 保存图片到本地
                            try:
                                with open(img_save_path, "wb") as f:
                                    f.write(image_part.blob)
                                # 记录图片信息（明确标注是文档中的组织架构图）
                                para_text.append(
                                    f"\n[图片{image_count}]：employee_manual.docx 中的公司组织架构图，保存路径：{img_save_path}\n")
                                image_count += 1
                            except Exception as e:
                                para_text.append(f"\n[图片{image_count}]：组织架构图提取失败（错误：{str(e)}）\n")

            # 合并段落内容（去除空值，确保文本连贯）
            para_content = "".join([p for p in para_text if p])
            if para_content:
                all_content.append(para_content)

        # ------------------- 处理表格（适配文档中8个规范表格，如部门职责、薪酬结构）-------------------
        elif child.tag == qn("w:tbl"):
            # 匹配当前表格对象（如技术部门职责、请假类型表格）
            table_idx = len(doc.tables) - len(doc.element.body.findall(qn("w:tbl"))) + doc.element.body.findall(
                qn("w:tbl")).index(child)
            table = doc.tables[table_idx]
            table_content = ["\n[表格开始]：employee_manual.docx 中的表格（如部门职责/薪酬等级）\n"]

            # 遍历表格行，保留表头与内容对应关系
            for row_idx, row in enumerate(table.rows):
                # 提取单元格文本（去除空白，适配文档表格的非空内容）
                row_cells = [cell.text.strip() for cell in row.cells if cell.text.strip()]
                if not row_cells:
                    continue  # 跳过空行

                # 转为 Markdown 格式（确保表格结构与原文档一致）
                markdown_row = f"| {' | '.join(row_cells)} |"
                table_content.append(markdown_row)

                # 表头下方添加分隔线（仅第一行后添加，提升可读性）
                if row_idx == 0:
                    separator = f"| {' | '.join(['---'] * len(row_cells))} |"
                    table_content.append(separator)

            # 标记表格结束
            table_content.append("[表格结束]\n")
            all_content.append("\n".join(table_content))

    # 4. 合并所有内容（段落、表格、图片信息），用空行分隔（符合文档阅读习惯）
    return "\n\n".join(all_content)


# ------------------- 测试：仅读取 employee_manual.docx -------------------
if __name__ == "__main__":
    # 确保文件路径正确（若文档不在当前目录，需替换为完整路径，如 "/Users/baimu/employee_manual.docx"）
    manual_content = load_employee_file("employee_manual.docx")

    # 将结果保存到 txt 文件（方便查看所有内容）
    with open("employee_manual_extracted.txt", "w", encoding="utf-8") as f:
        f.write(manual_content)

    print("✅ employee_manual.docx 读取成功！")
    print(f"📄 提取结果已保存到：employee_manual_extracted.txt")
    print(f"🖼️  图片（组织架构图）已保存到：./employee_manual_images/")
    print("\n🔍 内容预览（前800字符）：")
    print(manual_content[:800] + "...")